#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Query tools
#
# Software is free software released under the "GNU Affero General Public License v3.0"
#
# Copyright (c) 2014-2018  Alexandre Dulaunoy - a@foo.be
# Copyright (c) 2014-2018  Pieter-Jan Moreels - pieterjan.moreels@gmail.com

import os
import sys
import urllib.parse

import requests

runPath = os.path.dirname(os.path.realpath(__file__))
sys.path.append(os.path.join(runPath, ".."))

from lib.Toolkit import toStringFormattedCPE
from lib.CVEs import CveHandler
from lib.DatabaseLayer import findRanking, getCPE, getCVEs, cvesForCPE
from lib.Config import Configuration

rankinglookup = True
redisdb = Configuration.getRedisVendorConnection()

SCAN_COUNT = 1000


def findranking(cpe=None, loosy=True):
    i = None

    if cpe is None:
        return False

    result = False

    if loosy:
        for x in cpe.split(":"):
            if x != "":
                i = findRanking(cpe, regex=True)
            if i is None:
                continue
            if "rank" in i:
                result = i["rank"]
    else:
        i = findRanking(cpe, regex=True)
        print(cpe)
        if i is None:
            return result
        if "rank" in i:
            result = i["rank"]
    return result


def lookupcpe(cpeid=None):
    e = getCPE(cpeid)
    if e is None:
        return cpeid
    if "id" in e:
        return e["title"]


def lastentries(limit=5, namelookup=False, rankinglookup=True):
    entries = []
    for item in getCVEs(limit)["results"]:
        if not namelookup and rankinglookup is not True:
            entries.append(item)
        else:
            if "vulnerable_configuration" in item:
                vulconf = []
                ranking = []
                for conf in item["vulnerable_configuration"]:
                    if namelookup:
                        vulconf.append(lookupcpe(cpeid=conf))
                    else:
                        vulconf.append(conf)
                    if rankinglookup:
                        rank = findranking(cpe=conf)
                        if rank and rank not in ranking:
                            ranking.append(rank)
                item["vulnerable_configuration"] = vulconf
                if rankinglookup and len(ranking) > 0:
                    item["ranking"] = ranking
            entries.append(item)
    return entries


def apigetcve(api, cveid=None):
    if cveid is None:
        return False
    url = urllib.parse.urljoin(api, "api/cve/" + cveid)
    urltoget = urllib.parse.urljoin(url, cveid)

    with requests.Session() as session:
        r = session.get(urltoget)

    if r.status_code == 200:
        return r.text
    else:
        return False


def apibrowse(api, vendor=None, product=None):
    url = urllib.parse.urljoin(api, "api/browse")
    with requests.Session() as session:
        if vendor is None:
            r = session.get(url)
        elif product is None:
            urlvendor = url + "/" + vendor
            r = session.get(urlvendor)
        else:
            urlproduct = url + "/" + vendor + "/" + product
            r = session.get(urlproduct)

    if r.status_code == 200:
        return r.text
    else:
        return False


def apisearch(api, query=None):
    if query is None:
        return False
    url = urllib.parse.urljoin(api, "api/search/")
    url = url + query

    with requests.Session() as session:
        r = session.get(url)

    if r.status_code == 200:
        return r.text
    else:
        return False


# Lastly added
def qcvesForCPE(cpe, limit=0):
    cpe = toStringFormattedCPE(cpe)
    data = []
    if cpe:
        cvesp = CveHandler(
            rankinglookup=False, namelookup=False, via4lookup=True, capeclookup=False
        )
        r = cvesForCPE(cpe, limit=limit)
        for x in r["results"]:
            data.append(cvesp.getcve(x["id"]))
    return data


def getBrowseList(vendor):
    result = {}
    if (vendor is None) or type(vendor) == list:
        v1 = redisdb.smembers("o")
        v2 = redisdb.smembers("a")
        v3 = redisdb.smembers("h")
        vendor = sorted(list(set(list(v1) + list(v2) + list(v3))))
        cpe = None
    else:
        cpenum = redisdb.scard("v:" + vendor)
        if cpenum < 1:
            return None
        p = redisdb.smembers("v:" + vendor)
        cpe = sorted(list(p))
    result["vendor"] = vendor
    result["product"] = cpe
    return result


def getVersionsOfProduct(product):
    p = redisdb.smembers("p:" + product)
    return sorted(list(p))


def searchVendors(vendor_part):
    vendor_iterator = redisdb.scan_iter(f"v:*{vendor_part}*", count=SCAN_COUNT)
    vendor_list = []

    for vendor in vendor_iterator:
        vendor_list.append(vendor.replace("v:", ""))

    return {
        "vendor": vendor_list
    }


def searchProductsByVendor(vendor, product_part):
    product_iterator = redisdb.sscan_iter(f"v:{vendor}", f"*{product_part}*", count=SCAN_COUNT)
    product_list = []

    for product in product_iterator:
        product_list.append(product.replace("p:", ""))

    return {
        "product": product_list,
        "vendor": vendor
    }


def searchVersionsByProduct(vendor, product, version_part):
    version_iterator = redisdb.sscan_iter(f"p:{product}", f"*{version_part}*", count=SCAN_COUNT)
    version_list = list(version_iterator)

    return {
        "version": version_list,
        "product": product,
        "vendor": vendor
    }